library(coda)

### To use instead of hpd default
credible.interval <- function (x, p=0.95) {
  t(apply(x, 2, quantile, c((1-p)/2, 1 - (1-p)/2)))
}

### Grab contrasts for variables g1 and g2 for the names vars
contrast <- function (post, vars, g1, g2)
 {
     beta<-post$beta
     K <- post$K # n groups
     m.orig <- ncol(beta)/K
     g1.idx = which(colnames(post$gamma) == g1)
     g2.idx = which(colnames(post$gamma) == g2)

     vars.idx <- sapply(vars,
                        function (v) which(colnames(post$y)==v))

     ## Get a length(vars) list two-column matrices.  There are as
     ## many rows as samples in the posterior.  The first matrix holds
     ## g1 and g2 coefficients for vars[1], the second the g1 and g2
     ## coefficients for vars[2], and so on.
     sorted <- lapply(vars.idx,
                function (idx)
                      post$beta[, seq(idx, by=m.orig,
                                      length.out=K)[c(g1.idx, g2.idx)]])

     ## As many columns as vars, as many rows s samples
     return(sapply(sorted, function (m) m[,1] - m[,2]))
}

plotplist <- function (post, vars, vars.names=NULL,
                       base.group=colnames(post$gamma)[1],
                       other.groups=colnames(post$gamma)[-1],
                       group.names=NULL,
                       hpd.p=0.95,
                       #hpd.lty=c(2, 1), hpd.lwd=c(1, 2),
                       cex=1.5, 
                       interval=function (x, ...) HPDinterval(mcmc(x), ...),
                       return.contrasts=F,
                       ...)
{
    ## Basics
    beta <- post$beta
    K <- post$K # n groups
    m.orig <- ncol(beta) / K # n ind vars
    m <- length(vars)

    ## Set up display group names
    if (is.null(group.names))
        group.names=other.groups

    ## One matrix per other group.  Each matrix is samples of other
    ## group minus base group for each variable.
    contrasts <- lapply(other.groups, function (g)
                        contrast(post, vars, g, base.group))

    ### Means, and HPDS
    stats <-lapply(contrasts, function (c)
                   cbind(apply(c, 2, mean), interval(c, hpd.p)))
    y.max <- max(sapply(stats, max))
    y.min <- min(sapply(stats, min))

    pe <- as.vector(t(sapply(stats, function (x) x[,1])))
    lower <- as.vector(t(sapply(stats, function (x) x[,2])))
    upper <- as.vector(t(sapply(stats, function (x) x[,3])))
    plot(0, type='n', axes=F, xlab="", ylab="", xlim=c(0,length(pe)),
         ylim=c(y.min, y.max), ...)
    abline(h=0, lty=2, col="gray")
    abline(v=seq(K - 0.5, by=K-1, length.out=m-1), col="gray", lwd=3)
    text(pe, cex=cex, family="serif", labels=group.names)
    segments(1:length(lower), lower, 1:length(lower), upper)
    axis(1, labels=vars.names, at=seq(K/2, by=K-1, length.out=m), tick=F)
    ticks <- seq(min(lower), max(upper), length.out=5)
    axis(2, at=ticks, labels=round(ticks))
    if (return.contrasts)
      contrasts
}

### This is like plotplist, except it does every contrast for one IV
plotplist1 <- function (post, var, display.name=var,
                        groups=colnames(post$gamma),
                        group.names=groups,
                        hpd.p=0.95, cex=1.5,
                        mfrow=NULL,
                        interval=function (x, ...) HPDinterval(mcmc(x), ...),
                        print.title=T,
                        return.contrasts=F,
                        ...)
{
    if (is.null(mfrow))
        mfrow=c(length(groups), 1)

    mfrow.cur = par("mfrow")
    if (print.title)
      par(mfrow=mfrow, oma=c(0, 0, 2, 0))
    else
      par(mfrow=mfrow, oma=c(0, 0, 0, 0))
    contrasts <- list(length(groups))
    for (i in 1:length(groups)) {
        contrasts[[i]] <- plotplist(post, var,
                                    paste("Omitted Category =", group.names[i]),
                                    groups[i], groups[-i],
                                    group.names[-i], hpd.p, cex,
                                    interval,
                                    return.contrasts, ...)
    }
    if (print.title)
      title(display.name, outer=TRUE)
    par(mfrow=mfrow.cur)
    if (return.contrasts)
      contrasts
}

